package util

import (
	"log"
	"math"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync"
	"time"
)

// RPCClient returns rpc client
func RPCClient(network, addr string, timeout time.Duration) (*rpc.Client, error) {
	conn, err := net.DialTimeout(network, addr, timeout)
	if err != nil {
		return nil, err
	}
	return rpc.NewClient(conn), nil
}

// JSONRPCClient returns jsonrpc client
func JSONRPCClient(network, addr string, timeout time.Duration) (*rpc.Client, error) {
	conn, err := net.DialTimeout(network, addr, timeout)
	if err != nil {
		return nil, err
	}
	return jsonrpc.NewClient(conn), nil
}

// SingleConnRPCClient seal rpc client
type SingleConnRPCClient struct {
	sync.Mutex
	rpcClient *rpc.Client
	RPCServer string
	Timeout   time.Duration
}

// close rpc client
func (conn *SingleConnRPCClient) close() {
	if conn.rpcClient != nil {
		conn.rpcClient.Close()
		conn.rpcClient = nil
	}
}

func (conn *SingleConnRPCClient) serveConn() error {
	if conn.rpcClient != nil {
		return nil
	}

	var err error
	var retry = 1
	for {
		if conn.rpcClient != nil {
			return nil
		}
		// connect rpc server according to timeout and server address
		conn.rpcClient, err = JSONRPCClient("tcp", conn.RPCServer, conn.Timeout)
		if err != nil {
			log.Printf("dial %s fail: %v", conn.RPCServer, err)
			if retry > 3 {
				return err
			}
			time.Sleep(time.Duration(math.Pow(2.0, float64(retry))) * time.Second)
			retry++
			continue
		}
		return nil
	}
}

// Call rpc client
func (conn *SingleConnRPCClient) Call(method string, args interface{}, resp interface{}) error {
	// lock makes sure only one connection between agent and server
	conn.Lock()
	defer conn.Unlock()
	err := conn.serveConn()
	if err != nil {
		return err
	}

	timeout := time.Duration(10 * time.Second)
	done := make(chan error, 1)
	go func() {
		err := conn.rpcClient.Call(method, args, resp)
		done <- err
	}()
	select {
	case <-time.After(timeout):
		log.Printf("[WARN] rpc call timeout %v => %v", conn.rpcClient, conn.RPCServer)
		conn.close()
	case err := <-done:
		if err != nil {
			conn.close()
			return nil
		}
	}

	return nil
}
